"""
基于论文[1] Christian J. A. , Derksen H. , Watkins R. .Lunar crater identification in digital Images[J/OL].J. Astronaut. Sci.,2021,68(4):1056-1144

仅使用论文提供的不变量作为描述子，由于描述子没有顺序，因此只能使用投票法。
"""

import numpy as np
import torch
from utils.ellipse import radius_ellipse
from .base import TriadPyramid
from itertools import combinations


class ChristianPyramid(TriadPyramid):

    def triad_descriptor(
        self, i, j, k, Q1, Q2, Q3, *args, th=0.05, factor=0.3, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        N = Q3.shape[0]
        if len(Q1.shape) == 2:
            Q1 = Q1[None].repeat(N, dim=0)
        if len(Q2.shape) == 2:
            Q2 = Q2[None].repeat(N, dim=0)
        assert Q1.shape[0] == Q2.shape[0] == Q3.shape[0]
        d1 = np.mean(radius_ellipse(Q1), axis=0)
        d2 = np.mean(radius_ellipse(Q2), axis=0)
        d3 = np.mean(radius_ellipse(Q3), axis=0)
        # 按直径从小到大排序
        D = np.array((d1, d2, d3))
        # 任取两个直径计算比值，当比值接近1的数超过2个时，排除
        valid_ind = np.abs(D[[1, 2]] / D[0, None] - 1) > th
        valid_ind &= np.abs(D[[0, 2]] / D[1, None] - 1) > th
        valid_ind &= np.abs(D[[0, 1]] / D[2, None] - 1) > th
        valid_ind = valid_ind.all(axis=0)
        index = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[index, np.arange(N)]
        Q = torch.tensor(Q, device=self.device, dtype=torch.float32)
        # 计算不变量
        I1 = torch.einsum("nii->n", Q[0] @ torch.inverse(Q[1])) * torch.einsum(
            "nii->n", Q[1] @ torch.inverse(Q[0])
        )
        I2 = torch.einsum("nii->n", Q[1] @ torch.inverse(Q[2])) * torch.einsum(
            "nii->n", Q[2] @ torch.inverse(Q[1])
        )
        I3 = torch.einsum("nii->n", Q[2] @ torch.inverse(Q[0])) * torch.einsum(
            "nii->n", Q[0] @ torch.inverse(Q[2])
        )
        # 每行是一个3x6的点列，每三个元素代表一个点
        I = torch.stack((I1, I2, I3))
        # 按直径从小到大排序
        th = factor * I.abs()
        ijk = np.vstack((i, j, k))
        return ijk, I, th, np.argsort(index, axis=0), valid_ind

    def identify(self, params, uncertainty, factor, *args, max_try=5, **kwargs):
        """
        输入应该越接近于真实的检测输出越好。因此输入应当是一堆椭圆的拟合参数，输出应当是匹配确定的结果，即椭圆ID的列表。
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        # 将陨石坑按直径排序
        params = np.array(params)
        diameters = np.array(radius_ellipse(params)).mean(axis=0)
        idx = np.argsort(-diameters)
        # 从最大的直径开始，选取陨石坑构建金字塔
        if len(idx) < 4:
            return None
        fail_num = 0
        for i in range(idx.shape[0]):
            j, k, l = np.array(list(combinations(np.delete(idx, i), 3))).T
            i = idx[i].repeat(k.shape[0])
            # for k, l in combinations(not_used[2:], 2):
            # 批量算出金字塔的三元组
            ijk, I123, th, index, valid = self.triad_descriptor(
                np.concatenate((i, i, j, i)),
                np.concatenate((j, j, k, k)),
                np.concatenate((k, l, l, l)),
                params[np.concatenate((i, i, j, i))],
                params[np.concatenate((j, j, k, k))],
                params[np.concatenate((k, l, l, l))],
                factor,
            )
            # 变成四个三元组，即金字塔的结构
            valid = valid.reshape(4, -1).all(axis=0)
            ijk = ijk.reshape(3, 4, -1)[..., valid]
            I123 = I123.view(I123.shape[0], 4, -1)[..., valid]
            th = th.view(th.shape[0], 4, -1)[..., valid]
            index = index.reshape(3, 4, -1)[..., valid]
            if not valid.any():
                fail_num += 1
                continue
            fail_try = 0
            for it in range(ijk.shape[2]):
                ABC = self.triad_search(I123[..., it], index[..., it], th[..., it])
                if ABC is None:
                    continue
                pyra = ijk[..., it]
                assert np.unique(pyra[2, 1:]).shape[0] == 1
                ijkl = np.array((*pyra[:, 0], pyra[2, 1]))
                for IJKL in self.build_pyramids(ABC):
                    if len(IJKL) != 4:
                        continue
                    success, result = self.extend_identify(
                        params, uncertainty, ijkl, IJKL, *args, factor=factor, **kwargs
                    )
                    if success:
                        return result
                fail_try += 1
                if fail_try > max_try:
                    break
            fail_num += 1
            if fail_num > max_try:
                break
        return None
